Fix ModelParallel OOM issue during weight loading #21723
Conversation
amitsrivastava78
commented
Oct 7, 2025
amitsrivastava78
commented
- Modified load_own_variables() to use _direct_assign() for sharded variables
- Prevents loading full weight tensors on single device before distribution
- Resolves RESOURCE_EXHAUSTED errors when loading large models with ModelParallel
- Maintains backward compatibility for non-sharded variables
- Enables loading of models like Gemma2 2B/7B without OOM errors
- Added EinsumDense layer testing to ModelParallel sharded variable loading
- Modified load_own_variables() to use _direct_assign() for sharded variables - Prevents loading full weight tensors on single device before distribution - Resolves RESOURCE_EXHAUSTED errors when loading large models with ModelParallel - Maintains backward compatibility for non-sharded variables - Enables loading of models like Gemma2 2B/7B without OOM errors - Added EinsumDense layer testing to ModelParallel sharded variable loading
Summary of ChangesHello @amitsrivastava78, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical Out-Of-Memory (OOM) issue encountered when loading large models using ModelParallel in the JAX backend. The core problem stemmed from loading entire weight tensors onto a single device before distributing them, leading to Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request provides a crucial fix for an out-of-memory issue when loading sharded models with ModelParallel in the JAX backend. The core change, which involves using a new _direct_assign method to distribute weights before assigning them to device variables, is well-implemented and effectively prevents loading full tensors onto a single device. The addition of _ProtectedShardedArray and strong referencing to prevent premature garbage collection of sharded arrays is a clever solution to a common problem in JAX. The refactoring of sharding logic into a shared _initialize_variable_with_sharding helper improves code clarity and maintainability. The comprehensive end-to-end tests for sharded variable loading across various layer types provide strong confidence in the correctness of the fix. Overall, this is an excellent and thorough contribution that significantly improves the usability of model parallelism in Keras.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #21723 +/- ##
==========================================
- Coverage 82.59% 82.57% -0.02%
==========================================
Files 572 572
Lines 58535 58581 +46
Branches 9158 9166 +8
==========================================
+ Hits 48345 48375 +30
- Misses 7853 7862 +9
- Partials 2337 2344 +7
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- Fix PyTorch backend CI failures by adding _direct_assign method for proper numpy-to-tensor conversion - Restore JAX export functionality using jax_export.symbolic_shape for dynamic shape handling - Refactor variable loading logic to eliminate duplication between Dense and EinsumDense layers - Create shared utility function get_quantized_variable_load_order in keras/src/utils/variable_loading.py - Update layer implementations to use the shared variable loading utility - All tests passing: PyTorch backend, JAX backend, and layer-specific legacy loading tests
amitsrivastava78
force-pushed
the
pytorch-fix-rebased
branch
from
b0a7824 to
5da9108
Compare
- Improve host memory allocation for sharded variables by preferring JAX arrays over NumPy conversion - Remove unnecessary jax.block_until_ready() calls as JAX automatically blocks when needed - Add comprehensive documentation for memory stability protection and host allocation - Enhance logging for variable initialization and assignment operations - Add support for both NumPy and JAX arrays in variable assignment methods
There was a problem hiding this comment.
One overall note: We basically never log in the successful case, the logs are just way too noisy.
Can you create two separate small PRs:
- One for the fix for the initializer. My reading of it is that this is the only change needed:
def _initialize_with_initializer(self, initializer):
"""Initialize variable with initializer, running on CPU if sharding
is needed."""
if self._layout is not None:
# For sharded variables, run initializer on CPU to avoid device
# placement issues
with jax.default_device(jax.devices("cpu")[0]):
value = self._convert_to_tensor(
initializer(self._shape, dtype=self._dtype)
)
else:
# For non-sharded variables, use the default behavior
value = self._convert_to_tensor(
initializer(self._shape, dtype=self._dtype)
)
self._initialize(value)But we should check that it does what we think it does.
- One for the fix for the weight loading. It will have this change, and only this change, for the relevant layers:
for i, variable in enumerate(target_variables):
variable._direct_assign(store[str(i)])And jax_memory_cleanup should be addressed differently, which we can talk about.
- Remove _ProtectedShardedArray class and _maybe_create_strong_reference method from core.py - Remove jax.block_until_ready calls that are no longer needed - Simplify variable initialization and assignment logic - Remove all test cases related to reference holding from core_test.py - Tests now pass and are consistent with the simplified implementation
amitsrivastava78
force-pushed
the
pytorch-fix-rebased
branch
from
ec69373 to
92bf1ed
Compare
amitsrivastava78
force-pushed
the
pytorch-fix-rebased
branch
from
37f7475 to
92bf1ed
Compare
…1713 - Remove variable_loading.py (quantization/saving related) - Fix duplicate import in core_test.py - Revert layer files to remove quantization changes - Keep only core JAX memory management changes for OOM fix
- Remove get_quantized_variable_load_order imports from dense.py and einsum_dense.py - Replace function calls with inline variable ordering logic - Maintain compatibility with quantization loading
…n usage - Remove quantization-specific variable ordering in _legacy_load_own_variables - Keep _direct_assign usage for OOM prevention during sharded variable loading - Maintain compatibility with quantization_variable_spec
- Change dense.py and einsum_dense.py _legacy_load_own_variables to use _direct_assign - Maintains OOM prevention for ModelParallel while ensuring consistency across all layers - All layers now use _direct_assign for variable loading
There was a problem hiding this comment.
This looks good to me. But what about base_conv, embedding? Don't they need this change too?
https://github.com/search?q=repo%3Akeras-team%2Fkeras%20variable.assign(store%5Bstr(i)%5D)&type=code
Oh and what about base_optimizer?
google-ml-butler
bot
added
kokoro:force-run
ready to pull
- Change embedding.py, dense.py, and einsum_dense.py regular load_own_variables methods to use _direct_assign instead of assign - Ensures consistent OOM prevention for ModelParallel across all loading paths - base_conv.py and base_optimizer.py already used _direct_assign correctly - All variable loading now uses the same _direct_assign approach
google-ml-butler
bot
removed
the
ready to pull
Yes base_conv, embedding do have weights which needs to be sharded, same for base_optimizer |
…nt error handling
